import torch
from transformers import AutoModelForCausalLM
# from gemv_ext import LLMM1
from vllm import _custom_C

m = 4096
n = 1

for k in range(200, 1600, 8):
    with torch.no_grad():

        down_proj_weight = torch.rand(m, k, device="cuda", dtype=torch.float16)
        inp = torch.rand(n, k, device="cuda", dtype=torch.float16) - 0.5

        out_functional = torch.nn.functional.linear(inp, down_proj_weight)

        out_custom = torch.empty(n, m, dtype=inp.dtype, device="cuda")
        _custom_C.LLMM1(down_proj_weight, inp, out_custom, 4)

        absdiff = (out_functional.to(torch.float32) - out_custom.to(torch.float32)).abs()

        print("------ k:", k)
        print("  % off:", absdiff.mean().item() / out_functional.abs().mean().item())